# 获取邻接矩阵, 矩阵的值为权重

import pickle
import collections
import numpy as np
from bbdw.my_graph import Graph
import networkx as nx
from gensim.models import KeyedVectors


def node_2_index(file):
    with open(file, 'rb') as f:
        graph = pickle.load(f)
        nodes = graph.nodes
        print(len(nodes))
        word_nodes_len = 0
        doc_nodes_len = 0
        for node in nodes:
            if node.type == 'd':
                doc_nodes_len += 1
            else:
                word_nodes_len += 1

        doc_index = 0
        word_index = doc_nodes_len
        index_dict = {}
        for node in nodes:
            if node.type == 'd':
                index_dict[node.name] = doc_index
                doc_index += 1
            if node.type == 'w':
                index_dict[node.name] = word_index
                word_index += 1
    print('单词数：', word_nodes_len)
    print('文本数：', doc_nodes_len)
    with open(r'F:\mypython\final_subject\DWGI\data\index', 'wb') as f:
        pickle.dump(index_dict, f)

    # 一共11261个节点，其中0-4290为文档节点， 4291-11260为单词节点
    # 一共5782个节点，其中0-3971为文档节点， 3972-5781为单词节点


def get_adj(file):
    with open(file, 'rb') as f1:
        graph = pickle.load(f1)

    with open(r'F:\mypython\final_subject\DWGI\data\index', 'rb') as f2:
        index_dict = pickle.load(f2)

    edges = graph.edges
    nodes = graph.nodes
    adj = []
    for i in range(len(nodes)):
        temp = []
        for j in range(len(nodes)):
            temp.append(0)
        adj.append(temp)

    for edge in edges:
        node1_index = index_dict[edge.node1.name]
        node2_index = index_dict[edge.node2.name]

        adj[node1_index][node2_index] = edge.weight
        adj[node2_index][node1_index] = edge.weight
    adj = np.array(adj)
    print(adj)

    with open(r'F:\mypython\final_subject\DWGI\data\adj', 'wb') as f:
        pickle.dump(adj, f)


# 获得每个单词节点属于哪个子图 {0: 13, 1:34, ......单词编号:子图编号}
def divide_subgraph_index():
    def sub_graph(nodes, edges):
        with open(r'F:\mypython\final_subject\DWGI\data\index', 'rb') as f2:
            index_dict = pickle.load(f2)

        sub_graph_index = 0
        node_subgraph_dict = {}
        G = nx.Graph()
        for node1 in nodes:
            G.add_node(node1)
        for link in edges:
            G.add_edge(link[0], link[1])

        for c in nx.connected_components(G):
            # plot_graph(G.subgraph(c))
            nodeSet = G.subgraph(c).nodes
            edgeSet = G.subgraph(c).edges

            for node in nodeSet:
                node_subgraph_dict[index_dict[node]] = sub_graph_index
            sub_graph_index += 1

        return node_subgraph_dict


    word_edge = []
    word_list = []

    with open(r'F:\mypython\final_subject\bbdw\exclude_doc_single_word_out\connected_graph_exclude_doc.txt', 'rb') as f:
        graph = pickle.load(f)
        nodes_object = graph.nodes
        edges_object = graph.edges

        for edge in edges_object:
            word_edge.append([edge.node1.name, edge.node2.name])

        for node in nodes_object:
            word_list.append(node.name)

    point_list = word_list
    link_list = word_edge

    node_subgraph_dict = sub_graph(point_list, link_list)
    print(node_subgraph_dict)

    with open(r'F:\mypython\final_subject\DWGI\data\node_subgraph', 'wb') as f:
        pickle.dump(node_subgraph_dict, f)


def get_features():
    wv = KeyedVectors.load(r'F:\model\ai-lab\ai_lab_model.wv', mmap='r')
    with open(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt', 'rb') as f:
        graph = pickle.load(f)

    with open(r'F:\mypython\final_subject\DWGI\data\index', 'rb') as f:
        node_index = pickle.load(f)

    nodes = graph.nodes
    features = np.random.rand(len(nodes), 200)

    for node in nodes:
        try:
            if node.type == 'w':
                # word_emb = np.array(wv[node.name])
                features[node_index[node.name]] = wv[node.name]
        except:
            continue

    for node in nodes:
        if node.type == 'd':
            sum = np.zeros(200)
            neighbor = node.neighbors
            for n in neighbor:
                try:
                    sum += wv[n]
                except KeyError:
                    continue
            sum = sum / len(neighbor)
            if not np.any(sum):  # 不含有一个非零元素，即如果全部为0
                sum = np.random.rand(200)
            features[node_index[node.name]] = sum

    print(type(features))
    print(features)
    print(features.shape)
    with open(r'F:\mypython\final_subject\DWGI\data\features', 'wb') as f:
        pickle.dump(features, f)


# 获得每个子图中有哪些节点 {0：[1,2,3], 1:[5,7,9],...子图编号:[单词编号列表]}
def subgraph_node():
    with open(r'F:\mypython\final_subject\DWGI\data\node_subgraph', 'rb') as f:
        node_subgraph_dict = pickle.load(f)
        subgraph_node_dict = {}
        for key, val in node_subgraph_dict.items():
            node_index = key
            subgraph_index = val
            if subgraph_node_dict.get(subgraph_index) is None:
                subgraph_node_dict[subgraph_index] = [node_index]
            else:
                subgraph_node_dict[subgraph_index].append(node_index)
        print(subgraph_node_dict)

    with open(r'F:\mypython\final_subject\DWGI\data\subgraph_node', 'wb') as f:
        pickle.dump(subgraph_node_dict, f)


def generate():
    node_2_index(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt')
    get_adj(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt')
    divide_subgraph_index()
    get_features()
    subgraph_node()


if __name__ == '__main__':
    node_2_index(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt')
    get_adj(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt')
    divide_subgraph_index()
    get_features()
    subgraph_node()


# 处理单词不存在词表中的异常